feat(trainer): async checkpointing with pinned staging (supersedes #4331) - #4428

Async trainer checkpoint saves for every model.offload mode. This supersedes #4331 and builds on it.

What it does

ckpt.async_save = true (--ckpt.async-save for sft; --ckpt.async-save or --trainer.ckpt.async-save for rl; default off). A save copies each rank's model and optimizer state into pinned host buffers and returns. A background thread writes the copy with DCP while training continues.

Design

• Staging (PinnedStager). Each tensor in the state dict is copied into a pinned host buffer. Buffers are keyed by the tensor's path and reused across saves, so pinning happens on the first save only. Later staging calls take about 130 ms. DTensors are staged through their local shard and rewrapped with their spec. DCP's DefaultStager cannot stage this trainer's DTensors ("invalid python storage", also hit in feat(trainer): async DCP checkpoint saves #3874). Non-tensor values are deep-copied.
• Without offload (model.offload = "none"). The GPU→host copy runs on a side stream, overlapping the next step's forward and backward. The source tensors are kept alive until it finishes.
  • Fence: ckpt_manager.wait_for_staging() right before optimizer.step() in the RL and SFT loops. In every run here it waited 0 ms.
• With offload (optimizer / full). The state dict reads the CPU optimizer state where it lives. Unlike a synchronous save, it skips prepare_checkpoint_save's CPU→GPU→CPU round trip; this is feat(trainer): async checkpoint saves with CPU optimizer offload #4331's _build_state_dict path, including its GPU detour for fused entries split along their sharded dim. The CPU state is then copied into the pinned buffers before save returns.
  • It must not be passed by reference: offloaded state changes outside the optimizer step (full updates it during backward), and the copy means a write longer than one step never stalls the next step.
• Writer. A single-thread executor runs dcp.save on a dedicated gloo group, created on the first async save. Each write is submitted with its staging object as an argument. get_dcp_group() is cached and also serves the main thread's barriers and loads.
  • The thread is bound to the rank's GPU. A thread's current device starts at 0, so a writer thread created a CUDA context on GPU 0 for every rank: +3.7 GiB of non-PyTorch memory on GPU 0. feat(trainer): async checkpoint saves with CPU optimizer offload #4331's DCP thread writer has the same effect.
  • I chose a thread over DCP's writer process because the process path on torch 2.13/2.14 needs two workarounds: pickling a DTensor to the writer process segfaults, and staging needs /dev/shm the size of the shard.
  • The writer thread does not slow down steps without a save: 24.90 ± 0.11 s vs 24.88 ± 0.06 s without checkpointing.
• One save at a time. save first waits for the previous write. keep_last/keep_interval count a step only once its write has finished (credit feat(trainer): async checkpoint saves with CPU optimizer offload #4331), and at startup count only complete steps (see Completeness), so cleanup never deletes the newest finished checkpoint for a write in flight or for a step cut off mid-write before a resume.
• Errors surface at the fence, the next save, or the final wait(). The final save waits for its write before exit. The control API's checkpoint action now waits before mark_saved, so a committed checkpoint is on storage.
• Completeness (also fixes a bug in sync saves). DCP writes trainer/.metadata after every rank's shards. resolve_latest_ckpt_step now skips a step whose trainer checkpoint has none. Before this, a crash mid-write made --resume pick the partial step and fail in dcp_load (reproduced on feat(trainer): async checkpoint saves with CPU optimizer offload #4331 below). The trainer, the orchestrator and the launchers all resolve through this function, so they agree on the step.

How it differs from #4331

#4331 passes the CPU optimizer state to the writer by reference and holds the next optimizer update until the write finishes. When a write takes longer than a step, the step after each save stalls for the rest of the write. With the settings below, that is ~24 s per save, half of a sync save. It also only supports CPU optimizer offload, and it has no resume-side completeness check.

This PR keeps #4331's fence placement and keep_last handling, and its offloaded-state read path. It adds:

• pinned staging of the CPU state, so writes no longer stall training;
• offload = "none" support, with the overlapped copy and fence;
• the completeness check;
• a dedicated writer group;
• the control-action wait.

Cost

• Host: pinned buffers ≈ the rank's checkpoint shard (model + optimizer), kept for the run. With optimizer offload this comes on top of the offloaded state.
  • Measured growth in node used RAM, relative to the run without checkpointing, for a 20 GB checkpoint across 8 ranks:
    • offload optimizer: sync +14 GiB, feat(trainer): async checkpoint saves with CPU optimizer offload #4331 +33 GiB, this PR +49 GiB
    • offload none: sync +27 GiB, this PR +46–47 GiB
• GPU: with the device-bound writer, the peak is the same as sync: 21.4 GiB vs 21.4 GiB (offload none). PyTorch's allocated and reserved memory are unchanged.
  • Before the fix, peak memory rose 3.7 GiB (to 25.1, or 23.7 with offload optimizer, where feat(trainer): async checkpoint saves with CPU optimizer offload #4331 shows 23.6). The cause was the extra CUDA contexts on GPU 0, not staging buffers or allocator fragmentation.

Results

The equality, resume, kill and timing runs below used the first commit, where DCP's thread writer did the writes. After switching to the device-bound writer thread, I re-ran MoE equality (both offload modes, byte-identical), memory, and offload-none timing (save step 26.35 ± 1.91 s, step after save 25.1 s, peak GPU 21.4 GiB).

1 node, 8×H200, SFT trainer with fake data, checkpoints on node-local disk (~0.4 GB/s, so a 20 GB write takes ~48 s).

Equality. A synchronous reference copy of the same state was saved right before the async save, then compared chunk by chunk:

Resume. Resuming from the async step-4 checkpoint and from the sync copy gives losses identical to the uninterrupted run at steps 5–8:

• offload none: 8.9624 / 7.6331 / 7.9907 / 9.3335
• offload optimizer: 8.9665 / 7.6377 / 7.9933 / 9.3261

Kill mid-write. Async saves every step. Every trainer process was SIGKILLed while step_4/trainer had no .metadata.

• This PR (both offload modes): the resume logs "Skipping checkpoint step 4 … incomplete", resumes from step 3, and reproduces steps 4–5 (7.0710 / 8.3911).
• feat(trainer): async checkpoint saves with CPU optimizer offload #4331 in the same test: resolves step 4 and crashes (FileNotFoundError: …/step_4/trainer/.metadata).

Step time. Qwen3-1.7B, batch 256 (~25 s steps), ckpt.interval = 5, 16 steps, mean over saves at steps 5/10/15:

The first step after the first save is slower (28 s) while the pinned buffers are allocated and registered.

ruff check / ruff format are clean. tests/unit/utils/test_pathing.py passes, including the new test. tests/unit/test_configs.py has the same failures as on main (taskset packages missing in my environment). tests/unit/train/test_state_offload.py and tests/unit/train/test_ckpt.py pass on GPU (7 passed). Two of those tests fail on the code before the fix: test_offloaded_async_save_writes_after_it_returned (an offloaded async write that starts after save returned) and test_keep_last_ignores_a_step_cut_off_mid_write.

🤖 Generated with Claude Code